mega_moe: configurable SM headroom on cooperative mega grids (DG_MEGA_MOE_SM_HEADROOM) - #13

Summary

Reserve SM headroom on all cooperative mega-MoE grids via a shared helper, so a co-resident comm
kernel cannot starve the grid's residency and trip the software whole-grid barrier.

Credit: the forward-path fix originates from Ying Zhang's yingz/mega-sm-headroom branch
(commit 95046f0, opened here as #14). That branch was never merged, so the fix is absent from
affab4d — the commit shipped in the production trainer image. This PR re-applies it on the
current base, makes the amount configurable, and extends it to the launches that branch did not
cover. Comment text is adapted from Ying's original.

The defect

The SM100 FP8/FP4 and BF16 mega-MoE kernels (forward + backward) launch cooperative grids of
exactly kNumSMs CTAs (one per SM) joined by grid_sync / nvlink_barrier
(comm/barrier.cuh) with a hard 60 s timeout. The barrier is only satisfiable if every CTA is
simultaneously resident — but an ordinary launch gives no gang-scheduling guarantee: CUDA places
CTAs greedily, so a concurrent kernel holding some (not all) SMs leaves part of the grid queued
while resident CTAs spin against the deadline and never yield their slots. Mutual deadlock; at
60 s they trap, taking down the process (SIGABRT) and the GPU (Xid 43 fallout).

Evidence: INC-1291 (production, 2026-09-08)

Figma RLOR job xy98ftpyg7kovb4m (DSV4-Flash mxfp4, megamoe_precision=bf16, CP=4 EP=4 with
HSDP dp_replicate=2, 8× B300), mid-run forward_backward call_id=392:

DeepGEMM NVLink barrier timeout: rank=0, counter=286245, signal=3, target=4, phase=0, sign=0, tag=1
DeepGEMM grid sync timeout: sm=96, thread=0, grid_sync_idx=0, old=2147483710, current=2147483795, ...
Assertion failed: .../deep_gemm/comm/barrier.cuh:39, condition: false and "Grid sync timeout"
Assertion failed: .../deep_gemm/comm/barrier.cuh:80, condition: false and "NVLink barrier timeout"

Production-log decode: tag=1 = the forward kernel's dispatch barrier (kBeforeDispatchPullBarrierTag,
sm100_bf16_mega_moe.cuh:287); all 405 grid-sync prints at grid_sync_idx=0 (first barrier, before
any MMA); signal=3, target=4 = one EP rank's NVLink signal never arrived (its grid could not
become resident); one rank's counters show it was 38 CTAs short. 907 crash lines across all 8
ranks — full EP-group cascade, then Xid 43 + NCCL watchdog + SIGABRT. Not hardware: DCGM Xid was
0 on all 8 GPUs until the assert flipped it to 43, and the identical signature was reproduced
locally on healthy B300s with the production image.

Reproduced with real collectives on the production BF16 kernel (8× B300, the incident job's
shape): launching any FSDP-shaped collective — reduce-scatter, all-gather, or HSDP all-reduce —
immediately before the persistent kernel aborts it with the verbatim production signature. A
single isolated collective does not abort; back-to-back queued work does, because the deadlock
needs the SM-occupancy window to cover the launch — exactly what FSDP2's overlap streams produce
at every layer boundary. efsdp = dp_shard*cp*tp/(etp*ep) = 1 in this job shape, so the eFSDP
expert-grad RS is not the competitor; it is FSDP2's own AG/RS/AR.

Full RCA: fw-ai/fireworks#47752 (docs/incidents/2026-09-08-inc1291-deepgemm-grid-sync-timeout.md).

Changes

get_mega_moe_num_sms() in csrc/jit_kernels/impls/runtime_utils.hpp:
num_sms = get_num_sms() - DG_MEGA_MOE_SM_HEADROOM. The headroom is controlled exclusively
by the env var and has no built-in default — unset means 0, the historical full-device grid.
Deployments that run comm kernels concurrently with the mega grid must set it (sizing guidance in
the helper comment and the RCA). Non-negative and bounds asserts, rounded even because these are
2-CTA cluster launches (several sites assert num_sms % 2 == 0). Used at 6 launch sites:

• sm100_fp8_fp4_mega_moe.hpp — forward (Ying's original site)
• sm100_fp8_fp4_mega_moe_backward.hpp — all 4 backward launches (dgrad / combine / prelude / bf16-dgrad)
• sm100_bf16_mega_moe.hpp — default of the DG_BF16_MEGA_MOE_NUM_SMS override (override retained)

No kernel-side changes: kNumSMs is a template parameter fed from launch_args.grid_dim.first,
so the barrier arithmetic follows the host-side grid automatically.

The bf16 w2/w13 wgrads (sm100_bf16_mega_moe_wgrad.hpp) are ordinary GEMMs, not cooperative
grids — correctly excluded.

Why no default: the right value is deployment-specific (serving SendRecv: 2, Ying's
measured value; training FSDP2 collectives with 32 NCCL channels: ~8 SMs touched — one channel =
one 512-thread block, 4 blocks/SM). Baking either into the code silently misconfigures the other
deployment, so the env var is the single source of truth and unset means the historical
full-device grid. See the validation section for the measured training-side setting.

Known gap: the launch has no gang-scheduling contract

The underlying mismatch. This kernel semantically requires gang scheduling — either all
CTAs resident or none started — because the whole-grid barrier is only satisfiable under full
co-residency. But it uses an ordinary launch, and CUDA schedules CTAs greedily and
incrementally. The kernel needs a guarantee it never asked the driver for.

That is what makes the failure mode counterintuitive, and why reserving SMs can only ever be
partial mitigation. Measured locally (2× B300, production image):

So the grid does wait, and waiting is harmless when the whole grid waits. The kernel is only
fatal in the partial state — precisely the state greedy CTA scheduling produces whenever a
competitor holds some but not all SMs.

Multi-rank makes patience useless. Even if one rank's grid is fully queued (locally benign),
its peers are already resident and spinning in nvlink_barrier against their own independent
60 s deadlines. Any rank delayed past 60 s kills every other rank regardless of local residency —
the INC-1291 shape: one rank at barrier.cuh:39, peers at barrier.cuh:80.

Scope of this PR. Reserving SMs narrows the window in which a competitor can create the
partial state — it does not remove it. It does not cover a preemptive long-lived kernel
occupying SMs before the mega launch: verified locally, the abort still reproduces with the hog
launched ahead of the grid, with as little as one hogged SM, because cluster pairing means an
occupied SM also wastes its partner (enough free SMs != enough pairable SM pairs).

A structural fix needs one of:

• cudaLaunchCooperativeKernel — the driver validates occupancy up front and returns an error
  if the grid cannot be co-resident: fail fast at launch instead of trapping 60 s later and
  killing a multi-hour training job
• residency-agnostic synchronization — a barrier that does not assume simultaneous residency
  (work-queue / persistent-block style where CTAs can retire and be replaced)
• failing that, at minimum: make the timeout a recoverable error rather than a device-side
  trap, so it surfaces as a CUDA error the framework can retry instead of Xid 43 + SIGABRT

Filing this as the low-risk mitigation matching the prior serving fix; the structural fix should
be tracked separately.

Validation

• Production signature reproduced verbatim on stock affab4d (both timeout printfs, both
  asserts, SIGABRT, host-kernel Xid 43 line) — scripts in fw-ai/fireworks#47752
• Production BF16 kernel + real funcol collectives, incident job shape (8 GPU): all three
  collective types abort the stock build; zero aborts with DG_MEGA_MOE_SM_HEADROOM=8 +
  NCCL_MAX_NCHANNELS=8 + NCCL_MIN_NCHANNELS=8 + NCCL_NVLS_ENABLE=0 (90 iters mixed + 200
  iters e2e), with no measurable slowdown — collectives were ~9% faster at 64 MiB payloads,
  mega-moe within noise
• DG_MEGA_MOE_SM_HEADROOM verified honored: unset runs the full-device grid (0), value 200
  trips the bounds assert, value 8 runs
• Sizing reference: serving (SendRecv, few SMs) — 2 (Ying's measurement); the incident's
  training shape — 8 alongside NCCL_MAX_NCHANNELS=8/NCCL_MIN_NCHANNELS=8/NCCL_NVLS_ENABLE=0
• Trainer image with this commit (fd84a35) built and smoke-tested; shape-CI speed comparison
  on deepseek-v4-flash-0731-256k-lora in flight (fw-ai/fireworks#47752 will carry results)

🤖 Generated with Claude Code